package com.iwuyc.tools.commons.util.collection;

import com.google.common.collect.HashBasedTable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * jdk stream的工具类，用于简化stream复杂的操作。
 *
 * @author 吴宇春
 * @since 2025-03-11
 */
public class StreamUtils {
    /**
     * 覆盖替换，流式操作中，将对象收集为map类型的时候（{@link java.util.stream.Collectors#toMap(Function, Function)}）方法，出现重复键会抛异常的问题，在某些业务场景下，可以使用后面的重复键的对象覆盖前面的情况
     */
    public static final BinaryOperator<?> COVER_REPLACE = (old, newly) -> newly;
    /**
     * 保留最先放入到容器中的数据，丢弃新的数据
     */
    public static final BinaryOperator<?> DISCARD_NEWLY = (old, newly) -> old;

    /**
     * 从集合中提取出Set
     *
     * @param source 原集合
     * @param mapper 映射函数
     * @param <S>    原集合类型
     * @param <T>    映射类型
     * @return 映射后的Set
     */
    public static <S, T> Set<T> extractSet(Collection<S> source, Function<S, T> mapper) {
        return source.stream().map(mapper).collect(Collectors.toSet());
    }

    /**
     * 提取集合中元素的集合字段，平铺到set集合中，如果元素为null或者为空，则不添加到结果集合中
     *
     * @param sourceColl 待提取的集合
     * @param func       提取元素的函数
     * @param <S>        源类型
     * @param <T>        目标类型
     * @return 提取后的结果集合
     */
    public static <S, T> Set<T> extractCollToSet(Collection<S> sourceColl, Function<S, Collection<T>> func) {
        if (CollectionUtil.isEmpty(sourceColl)) {
            return Collections.emptySet();
        }
        final Set<T> result = new HashSet<>();
        for (S item : sourceColl) {
            final Collection<T> valColl = func.apply(item);
            if (CollectionUtil.isEmpty(valColl)) {
                continue;
            }
            result.addAll(valColl);
        }
        return result;
    }

    public static <S, K, V> Map<K, V> toMap(Collection<S> source, Function<S, K> keyMapper, Function<S, V> valueMapper) {
        return toMap(source, keyMapper, valueMapper, coverReplace(), HashMap::new);
    }

    public static <S, K, V> Map<K, V> toMap(Collection<S> source, Function<S, K> keyMapper, Function<S, V> valueMapper, BinaryOperator<V> mergeFunction, Supplier<Map<K, V>> mapSupplier) {
        return source.stream().collect(Collectors.toMap(keyMapper, valueMapper, mergeFunction, mapSupplier));
    }


    /**
     * 覆盖替换，流式操作中，将对象收集为map类型的时候（{@link java.util.stream.Collectors#toMap(Function, Function)}）方法，出现重复键会抛异常的问题，在某些业务场景下，可以使用后面的重复键的对象覆盖前面的情况
     *
     * @param <T> 容器中的对象类型
     * @return 覆盖式替换的操作类
     * @author 吴宇春
     */
    @SuppressWarnings("unchecked")
    public static <T> BinaryOperator<T> coverReplace() {
        return (BinaryOperator<T>) COVER_REPLACE;
    }

    /**
     * 保留最先放入到容器中的数据，丢弃新的数据
     *
     * @param <T> 容器中的对象类型
     * @return 操作类实例
     * @author 吴宇春
     */
    @SuppressWarnings("unchecked")
    public static <T> BinaryOperator<T> discardNewly() {
        return (BinaryOperator<T>) DISCARD_NEWLY;
    }

    /**
     * 将sources中的对象，按照指定的数值，进行分组，返回分组后的结果
     *
     * @param sources   源数据
     * @param keyMapper 分组key的映射函数
     * @param <S>       源数据类型
     * @param <K>       分组key的类型
     * @return 分组后的数据
     * @author Neil
     * @since 2025-03-12
     */
    public static <S, K> Map<K, Collection<S>> groupBy(Collection<S> sources, Function<S, K> keyMapper) {
        return groupBy(sources, keyMapper, Function.identity());
    }

    /**
     * 将sources中的对象，按照指定的数值，进行分组，返回分组后的结果
     *
     * @param sources     源数据
     * @param keyMapper   分组key的映射函数
     * @param valueMapper 分组value的映射函数
     * @param <S>         源数据类型
     * @param <K>         分组key的类型
     * @param <V>         分组value的类型
     * @return 分组后的数据
     * @author Neil
     * @since 2025-03-12
     */
    public static <S, K, V> Map<K, Collection<V>> groupBy(Collection<S> sources, Function<S, K> keyMapper, Function<S, V> valueMapper) {
        if (CollectionUtil.isEmpty(sources)) {
            return Collections.emptyMap();
        }
        Map<K, Collection<V>> result = new HashMap<>(sources.size());
        for (S item : sources) {
            K key = keyMapper.apply(item);
            final V val = valueMapper.apply(item);
            final Collection<V> valColl = result.computeIfAbsent(key, k -> new ArrayList<>());
            valColl.add(val);
        }
        return result;
    }


    /**
     * 计算源数据中，映射为指定类型的数据，去重后的数量
     *
     * @param sources 源数据列表
     * @param <S>     源数据类型
     * @return 去重后的数量
     */
    public static <S> long counterDistinct(Collection<S> sources) {
        return counterDistinct(sources, Function.identity());
    }

    /**
     * 计算源数据中，映射为指定类型的数据，去重后的数量
     *
     * @param sources 源数据列表
     * @param mapper  映射函数
     * @param <S>     源数据类型
     * @return 去重后的数量
     */
    public static <S> long counterDistinct(Collection<S> sources, Function<S, ?> mapper) {
        return sources.stream().map(mapper).distinct().count();
    }


    /**
     * 将集合转换为HashBasedTable行，cell 数据为 SOURCE 实例
     *
     * @param sources      源数据集合
     * @param rowGetter    行获取器
     * @param columnGetter 列获取器
     * @param <SOURCE>     源数据类型
     * @param <ROW>        行类型
     * @param <COLUMN>     列类型
     * @return HashBasedTable 转换后的实例
     * @see HashBasedTable
     */
    public static <SOURCE, ROW, COLUMN> HashBasedTable<ROW, COLUMN, SOURCE> toHashBaseTable(
            Collection<SOURCE> sources,
            Function<SOURCE, ROW> rowGetter,
            Function<SOURCE, COLUMN> columnGetter
    ) {
        return toHashBaseTable(sources, rowGetter, columnGetter, Function.identity());
    }

    /**
     * 将集合转换为HashBasedTable行
     *
     * @param sources        源数据集合
     * @param rowGetter      行获取器
     * @param columnGetter   列获取器
     * @param cellDataGetter 单元格数据获取器
     * @param <SOURCE>       源数据类型
     * @param <ROW>          行类型
     * @param <COLUMN>       列类型
     * @param <VALUE>        单元格数据类型
     * @return HashBasedTable 转换后的实例
     * @see HashBasedTable
     */
    public static <SOURCE, ROW, COLUMN, VALUE> HashBasedTable<ROW, COLUMN, VALUE> toHashBaseTable(
            Collection<SOURCE> sources,
            Function<SOURCE, ROW> rowGetter,
            Function<SOURCE, COLUMN> columnGetter,
            Function<SOURCE, VALUE> cellDataGetter
    ) {
        return toHashBaseTable(sources, rowGetter, columnGetter, cellDataGetter, false);
    }

    /**
     * 将集合转换为HashBasedTable行
     *
     * @param sources        源数据集合
     * @param rowGetter      行获取器
     * @param columnGetter   列获取器
     * @param cellDataGetter 单元格数据获取器
     * @param ignoreNull     是否过滤value为null的数据
     * @param <SOURCE>       源数据类型
     * @param <ROW>          行类型
     * @param <COLUMN>       列类型
     * @param <VALUE>        单元格数据类型
     * @return HashBasedTable 转换后的实例
     * @see HashBasedTable
     */
    public static <SOURCE, ROW, COLUMN, VALUE> HashBasedTable<ROW, COLUMN, VALUE> toHashBaseTable(
            Collection<SOURCE> sources,
            Function<SOURCE, ROW> rowGetter,
            Function<SOURCE, COLUMN> columnGetter,
            Function<SOURCE, VALUE> cellDataGetter,
            boolean ignoreNull
    ) {
        final HashBasedTable<ROW, COLUMN, VALUE> result = HashBasedTable.create();
        if (CollectionUtil.isEmpty(sources)) {
            return result;
        }

        for (SOURCE item : sources) {
            final VALUE cellData = cellDataGetter.apply(item);
            if (ignoreNull && cellData == null) {
                continue;
            }
            result.put(rowGetter.apply(item), columnGetter.apply(item), cellData);
        }
        return result;
    }
}
